package rocksdb

import (
	"encoding/hex"
	"fmt"

	"github.com/linxGnu/grocksdb"
)

type RocksDBService struct {
	numShards    int
	shardManager *ShardManager
}

func NewRocksDBService(numShards int, basePath string) (*RocksDBService, error) {
	sm, err := NewShardManager(numShards, basePath)
	if err != nil {
		return nil, err
	}
	return &RocksDBService{
		numShards:    numShards,
		shardManager: sm,
	}, nil
}

func (qs *RocksDBService) Query(key []byte) (string, error) {
	db, err := qs.shardManager.GetShard(key)
	if err != nil {
		return "", err
	}
	ro := grocksdb.NewDefaultReadOptions()
	ro.SetFillCache(true)
	defer ro.Destroy()
	data, err := db.Get(ro, []byte(key))
	if err != nil {
		return "", err
	}
	defer data.Free()
	if !data.Exists() {
		return "", fmt.Errorf("key not found")
	}
	return string(data.Data()), nil
}

func (qs *RocksDBService) BatchQuery(md5Hashes [][]byte) (map[string]string, error) {
	result := make(map[string]string)
	for _, hash := range md5Hashes {
		value, err := qs.Query(hash)
		if err != nil {
			result[hex.EncodeToString(hash)] = "" // 或者记录错误
			continue
		}
		result[hex.EncodeToString(hash)] = value
	}
	return result, nil
}

func (qs *RocksDBService) Write(key, value []byte) error {
	db, err := qs.shardManager.GetShard(key)
	if err != nil {
		return err
	}
	wr := grocksdb.NewDefaultWriteOptions()
	defer wr.Destroy()
	err = db.Put(wr, key, value)
	db.Write(wr, &grocksdb.WriteBatch{})
	if err != nil {
		return err
	}
	return nil
}

func (qs *RocksDBService) Delete(key []byte) error {
	db, err := qs.shardManager.GetShard(key)
	if err != nil {
		return err
	}
	db.Delete(grocksdb.NewDefaultWriteOptions(), key)
	return nil
}
